package com.ruoyi.flink.source;

import com.ruoyi.flink.bean.TaxiFare;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;

/**
 * @program: ruoyi
 * @description: MyTaxiFareSource
 * @author: zengchen
 * @create: 2020-11-28 16:41
 **/
public class MyTaxiFareSource extends RichSourceFunction<TaxiFare> implements CheckpointedFunction {

    private volatile boolean running = true;

    private transient ListState<Long> checkpointedCount;

    private long count = 0L;

    @Override
    public void run(SourceContext<TaxiFare> ctx) throws Exception {
        long count = 1;
        while (running && count < 1000) {
            synchronized (ctx.getCheckpointLock()) {
                TaxiFare fare = new TaxiFare(count);
                ctx.collectWithTimestamp(fare, fare.getEventTime());
//                ctx.emitWatermark();
//                Thread.sleep(2000);
                count++;
            }
        }
    }

    @Override
    public void cancel() {
        running = false;
    }

    /**
     * This method is called when a snapshot for a checkpoint is requested. This acts as a hook to the function to
     * ensure that all state is exposed by means previously offered through {@link FunctionInitializationContext} when
     * the Function was initialized, or offered now by {@link FunctionSnapshotContext} itself.
     *
     * @param context the context for drawing a snapshot of the operator
     * @throws Exception Thrown, if state could not be created ot restored.
     */
    @Override
    public void snapshotState(FunctionSnapshotContext context) throws Exception {
        this.checkpointedCount.clear();
        this.checkpointedCount.add(count);
    }

    /**
     * This method is called when the parallel function instance is created during distributed
     * execution. Functions typically set up their state storing data structures in this method.
     *
     * @param context the context for initializing the operator
     * @throws Exception Thrown, if state could not be created ot restored.
     */
    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {
        this.checkpointedCount = context
                .getOperatorStateStore()
                .getListState(new ListStateDescriptor<>("count", Long.class));

        if (context.isRestored()) {
            for (Long count : this.checkpointedCount.get()) {
                this.count = count;
            }
        }
    }
}
